{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PnqQCtESuc8p"
      },
      "source": [
        "Data Pre-processing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_3wWiPQldzlG"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import math\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import TensorDataset, DataLoader\n",
        "from torch.optim.lr_scheduler import _LRScheduler\n",
        "\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm\n",
        "import sklearn\n",
        "from copy import deepcopy\n",
        "\n",
        "df = pd.read_csv(\"../aqi_city_data_v2_unrolled.csv\")\n",
        "\n",
        "DROP_ONEHOT = True\n",
        "SEQ_LENGTH = 7\n",
        "\n",
        "if DROP_ONEHOT:\n",
        "  INPUT_DIM = 10 \n",
        "else:\n",
        "  INPUT_DIM = 29\n",
        "\n",
        "HIDDEN_DIM = 32\n",
        "LAYER_DIM = 3\n",
        "\n",
        "normalization_type = 'mean_std' # 'max', mean_std\n",
        "\n",
        "def get_train_test_data(df):\n",
        "  # we'll mostly need median and variance values of features for most of our needs\n",
        "\n",
        "  for col in df.columns:\n",
        "    for x in [\"min\", \"max\", \"count\", \"County\", \"past_week\", \"latitude\", \"longitude\", \"State\", \"variance\"]:\n",
        "      if x in col:\n",
        "        df.drop([col], axis=1, inplace=True)\n",
        "\n",
        "  df[\"Population Staying at Home\"] = df[\"Population Staying at Home\"].apply(lambda x: x.replace(\",\", \"\"))\n",
        "  df[\"Population Not Staying at Home\"] = df[\"Population Not Staying at Home\"].apply(lambda x: x.replace(\",\", \"\"))\n",
        "\n",
        "  # Now we want 2 more features. Which day of week it is and which month it is.\n",
        "  # Both of these will be one-hot and hence we'll add 7+12 = 19 more columns.\n",
        "  # Getting month id is easy from the datetime column. \n",
        "  # For day of week, we'll use datetime library.\n",
        "  \n",
        "  df['weekday'] = df['Date'].apply(lambda x: datetime.strptime(x, \"%Y-%m-%d\").weekday())\n",
        "  df['month'] = df['Date'].apply(lambda x: datetime.strptime(x, \"%Y-%m-%d\").month - 1)\n",
        "\n",
        "  # using one-hot on month and weekday\n",
        "  weekday_onehot = pd.get_dummies(df['weekday'])\n",
        "  weekday_onehot.columns = [\"day_\"+str(x) for x in weekday_onehot]\n",
        "  month_onehot = pd.get_dummies(df['month'])\n",
        "  month_onehot.columns = [\"month_\"+str(x) for x in month_onehot]\n",
        "\n",
        "  df.drop(['weekday', 'month'], axis=1, inplace=True)\n",
        "  df = df.join([weekday_onehot, month_onehot])\n",
        "\n",
        "  cities_list = list(set(df['City']))\n",
        "  city_df = {}\n",
        "  test_indices_of_cities = {}\n",
        "  train_set = {}\n",
        "  test_set = {}\n",
        "  TEST_SET_SIZE = 60\n",
        "\n",
        "  for city in cities_list:\n",
        "    city_df[city] = df[df['City'] == city].sort_values('Date').reset_index()\n",
        "    for col in city_df[city].columns:\n",
        "      if col in [\"pm25_median\", \"o3_median\", \"so2_median\", \"no2_median\", \"pm10_median\", \"co_median\"]:\n",
        "        continue\n",
        "      try:  \n",
        "        _mean = np.nanmean(city_df[city][col])\n",
        "        if np.isnan(_mean) == True:\n",
        "          _mean = 0\n",
        "        city_df[city][col] = city_df[city][col].fillna(_mean)\n",
        "      except:\n",
        "        pass\n",
        "    \n",
        "    random.seed(0)\n",
        "    test_index_start = random.randint(0, city_df[city].shape[0] - TEST_SET_SIZE)\n",
        "    test_indices_of_cities[city] = [test_index_start, test_index_start + TEST_SET_SIZE]\n",
        "\n",
        "    test_set[city] = city_df[city].iloc[test_index_start:test_index_start + TEST_SET_SIZE]\n",
        "    train_set[city] = city_df[city].drop(index=list(range(test_index_start, test_index_start + TEST_SET_SIZE)))\n",
        "\n",
        "  return train_set, test_set\n",
        "\n",
        "train_set, test_set = get_train_test_data(df)\n",
        "\n",
        "cities_list = list(train_set.keys())\n",
        "\n",
        "all_train = pd.DataFrame()\n",
        "for city in cities_list:\n",
        "  all_train = all_train.append(train_set[city], ignore_index=True)\n",
        "\n",
        "all_test = pd.DataFrame({})\n",
        "for city in test_set:\n",
        "  all_test = all_test.append(test_set[city], ignore_index=True)\n",
        "\n",
        "concat_df = pd.concat([all_train,all_test],axis=0)\n",
        "\n",
        "# ---------------------------------------------------------------------------- #\n",
        "col_max = {}\n",
        "col_mean = {}\n",
        "col_mean2 = {}\n",
        "col_std = {}\n",
        "\n",
        "for city in cities_list:\n",
        "  col_mean[city] = {}\n",
        "  for col in train_set[city]:\n",
        "    if col in [\"index\", \"Date\", \"City\"]:\n",
        "      continue\n",
        "\n",
        "    train_set[city][col] = train_set[city][col].astype(\"float\")\n",
        "    test_set[city][col] = test_set[city][col].astype(\"float\")\n",
        "\n",
        "    if col in [\"pm25_median\", \"o3_median\", \"so2_median\", \"no2_median\", \"pm10_median\", \"co_median\"]:\n",
        "      _mean = np.nanmean(train_set[city][col])\n",
        "      if np.isnan(_mean) == True:\n",
        "        _mean = 0\n",
        "      \n",
        "      col_mean[city][col] = _mean\n",
        "      train_set[city][col] = train_set[city][col].fillna(_mean)\n",
        "\n",
        "    if normalization_type == 'mean_std':\n",
        "      col_mean2[col] = np.nanmean(concat_df[col].astype(\"float\"))\n",
        "      col_std[col] = np.nanstd(concat_df[col].astype(\"float\"))\n",
        "      train_set[city][col] = (train_set[city][col] - col_mean2[col]) / (col_std[col] + 0.001)\n",
        "      test_set[city][col] = (test_set[city][col] - col_mean2[col]) / (col_std[col] + 0.001)\n",
        "\n",
        "    else:\n",
        "      col_max[col] = concat_df[col].astype(\"float\").max()\n",
        "      train_set[city][col] = train_set[city][col] / (col_max[col] + 0.001)\n",
        "      test_set[city][col] = test_set[city][col] / (col_max[col] + 0.001)\n",
        "\n",
        "  if DROP_ONEHOT:\n",
        "    train_set[city].drop(train_set[city].columns[-19:], axis=1, inplace=True)\n",
        "    test_set[city].drop(test_set[city].columns[-19:], axis=1, inplace=True)\n",
        "\n",
        "class CityDataP(torch.utils.data.Dataset):\n",
        "  def __init__(self, selected_column, split):\n",
        "    self.split = split\n",
        "    if split == \"train\":\n",
        "      self.dataset = train_set\n",
        "    else:\n",
        "      self.dataset = test_set\n",
        "\n",
        "    self.valid_city_idx = 0\n",
        "    self.valid_day_idx = 0\n",
        "    self.selected_column = selected_column\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    if self.split != \"train\":\n",
        "      # getting all data out of the validation set\n",
        "      out, city = self.get_idx_data(idx)\n",
        "    \n",
        "    else:\n",
        "      # getting data randomly for train split\n",
        "      city = random.choice(cities_list)\n",
        "      _df = self.dataset[city]\n",
        "      start_idx = random.randint(1,_df.shape[0]-SEQ_LENGTH)\n",
        "      out =  _df.iloc[start_idx-1:start_idx+SEQ_LENGTH]\n",
        "\n",
        "    out = out.drop(['index', 'Date', 'City'], axis=1)\n",
        "\n",
        "    Y = pd.DataFrame({})\n",
        "    Y_true = pd.DataFrame({})\n",
        "\n",
        "    for col in out.columns:\n",
        "      if col == self.selected_column:\n",
        "        Y_true[col] = out[col]\n",
        "        Y[col] = out[col].fillna(col_mean[city][col])\n",
        "      \n",
        "      if col in [\"pm25_median\", \"pm10_median\", \"o3_median\", \"so2_median\", \"no2_median\", \"co_median\"]:\n",
        "        out.drop([col], axis=1, inplace=True)\n",
        "      else:\n",
        "        out[col] = out[col].astype(\"float\")\n",
        "\n",
        "    out = np.concatenate((np.array(out)[1:,:], np.array(Y)[:-1,:]), axis=1)\n",
        "    Y = np.array(Y)[1:]\n",
        "    Y_true = np.array(Y_true)[1:]\n",
        "    \n",
        "    return out, Y, Y_true\n",
        "\n",
        "  def get_idx_data(self, idx):\n",
        "    city = cities_list[self.valid_city_idx]\n",
        "    _df = self.dataset[city]\n",
        "\n",
        "    out =  _df.iloc[self.valid_day_idx:self.valid_day_idx+SEQ_LENGTH]\n",
        "    \n",
        "    if self.valid_day_idx+SEQ_LENGTH >= _df.shape[0]:\n",
        "      self.valid_day_idx = 0\n",
        "      self.valid_city_idx += 1\n",
        "    else:\n",
        "      self.valid_day_idx += 1\n",
        "\n",
        "    return out, city\n",
        "\n",
        "  def __len__(self):\n",
        "    if self.split != \"train\":\n",
        "      return (61-SEQ_LENGTH)*len(cities_list)\n",
        "    return len(all_train) - (SEQ_LENGTH - 1)*len(cities_list)\n",
        "\n",
        "class CityDataForecast(torch.utils.data.Dataset):\n",
        "  def __init__(self, selected_column, split):\n",
        "    self.split = split\n",
        "    if split == \"train\":\n",
        "      self.dataset = train_set\n",
        "    else:\n",
        "      self.dataset = test_set\n",
        "\n",
        "    self.valid_city_idx = 0\n",
        "    self.valid_day_idx = 0\n",
        "    self.selected_column = selected_column\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    if self.split != \"train\":\n",
        "      # getting all data out of the validation set\n",
        "      out, city = self.get_idx_data(idx)\n",
        "    \n",
        "    else:\n",
        "      # getting data randomly for train split\n",
        "      city = random.choice(cities_list)\n",
        "      _df = self.dataset[city]\n",
        "      start_idx = random.randint(1,_df.shape[0]-SEQ_LENGTH)\n",
        "      out =  _df.iloc[start_idx-1:start_idx+SEQ_LENGTH]\n",
        "\n",
        "    out = out.drop(['index', 'Date', 'City'], axis=1)\n",
        "\n",
        "    Y = pd.DataFrame({})\n",
        "    Y_true = pd.DataFrame({})\n",
        "\n",
        "    for col in out.columns:\n",
        "      if col == self.selected_column:\n",
        "        Y_true[col] = out[col]\n",
        "        #print(out[col])\n",
        "        Y[col] = out[col].fillna(col_mean[city][col])\n",
        "      \n",
        "      if col in [\"pm25_median\", \"pm10_median\", \"o3_median\", \"so2_median\", \"no2_median\", \"co_median\"]:\n",
        "        out.drop([col], axis=1, inplace=True)\n",
        "      else:\n",
        "        out[col] = out[col].astype(\"float\")\n",
        "\n",
        "    out = np.concatenate((np.array(out)[1:,:], np.array(Y)[:-1,:]), axis=1)\n",
        "    Y = np.array(Y)[1:]\n",
        "\n",
        "    Y_true = np.array(Y_true)[1:]\n",
        "\n",
        "    return out, Y, Y_true\n",
        "\n",
        "  def get_idx_data(self, idx):\n",
        "    city = cities_list[self.valid_city_idx]\n",
        "    _df = self.dataset[city]\n",
        "\n",
        "    out =  _df.iloc[self.valid_day_idx:self.valid_day_idx+SEQ_LENGTH]\n",
        "    \n",
        "    if self.valid_day_idx+SEQ_LENGTH >= _df.shape[0]:\n",
        "      self.valid_day_idx = 0\n",
        "      self.valid_city_idx += 1\n",
        "    else:\n",
        "      self.valid_day_idx += 1\n",
        "\n",
        "    return out, city\n",
        "\n",
        "  def __len__(self):\n",
        "    if self.split != \"train\":\n",
        "      return (61-SEQ_LENGTH)*len(cities_list)\n",
        "    return len(all_train) - (SEQ_LENGTH - 1)*len(cities_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaV-7ufAukFF"
      },
      "source": [
        "Loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KV3cLeL2ojvI"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "Code by Mehran Maghoumi\n",
        "link: https://github.com/Maghoumi/pytorch-softdtw-cuda/blob/master/soft_dtw_cuda.py\n",
        "'''\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.cuda\n",
        "from numba import jit\n",
        "from torch.autograd import Function\n",
        "from numba import cuda\n",
        "import math\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "@cuda.jit\n",
        "def compute_softdtw_cuda(D, gamma, bandwidth, max_i, max_j, n_passes, R):\n",
        "    \"\"\"\n",
        "    :param seq_len: The length of the sequence (both inputs are assumed to be of the same size)\n",
        "    :param n_passes: 2 * seq_len - 1 (The number of anti-diagonals)\n",
        "    \"\"\"\n",
        "    # Each block processes one pair of examples\n",
        "    b = cuda.blockIdx.x\n",
        "    # We have as many threads as seq_len, because the most number of threads we need\n",
        "    # is equal to the number of elements on the largest anti-diagonal\n",
        "    tid = cuda.threadIdx.x\n",
        "\n",
        "    # Compute I, J, the indices from [0, seq_len)\n",
        "\n",
        "    # The row index is always the same as tid\n",
        "    I = tid\n",
        "\n",
        "    inv_gamma = 1.0 / gamma\n",
        "\n",
        "    # Go over each anti-diagonal. Only process threads that fall on the current on the anti-diagonal\n",
        "    for p in range(n_passes):\n",
        "\n",
        "        # The index is actually 'p - tid' but need to force it in-bounds\n",
        "        J = max(0, min(p - tid, max_j - 1))\n",
        "\n",
        "        # For simplicity, we define i, j which start from 1 (offset from I, J)\n",
        "        i = I + 1\n",
        "        j = J + 1\n",
        "\n",
        "        # Only compute if element[i, j] is on the current anti-diagonal, and also is within bounds\n",
        "        if I + J == p and (I < max_i and J < max_j):\n",
        "            # Don't compute if outside bandwidth\n",
        "            if not (abs(i - j) > bandwidth > 0):\n",
        "                r0 = -R[b, i - 1, j - 1] * inv_gamma\n",
        "                r1 = -R[b, i - 1, j] * inv_gamma\n",
        "                r2 = -R[b, i, j - 1] * inv_gamma\n",
        "                rmax = max(max(r0, r1), r2)\n",
        "                rsum = math.exp(r0 - rmax) + math.exp(r1 - rmax) + math.exp(r2 - rmax)\n",
        "                softmin = -gamma * (math.log(rsum) + rmax)\n",
        "                R[b, i, j] = D[b, i - 1, j - 1] + softmin\n",
        "\n",
        "        # Wait for other threads in this block\n",
        "        cuda.syncthreads()\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "@cuda.jit\n",
        "def compute_softdtw_backward_cuda(D, R, inv_gamma, bandwidth, max_i, max_j, n_passes, E):\n",
        "    k = cuda.blockIdx.x\n",
        "    tid = cuda.threadIdx.x\n",
        "\n",
        "    # Indexing logic is the same as above, however, the anti-diagonal needs to\n",
        "    # progress backwards\n",
        "    I = tid\n",
        "\n",
        "    for p in range(n_passes):\n",
        "        # Reverse the order to make the loop go backward\n",
        "        rev_p = n_passes - p - 1\n",
        "\n",
        "        # convert tid to I, J, then i, j\n",
        "        J = max(0, min(rev_p - tid, max_j - 1))\n",
        "\n",
        "        i = I + 1\n",
        "        j = J + 1\n",
        "\n",
        "        # Only compute if element[i, j] is on the current anti-diagonal, and also is within bounds\n",
        "        if I + J == rev_p and (I < max_i and J < max_j):\n",
        "\n",
        "            if math.isinf(R[k, i, j]):\n",
        "                R[k, i, j] = -math.inf\n",
        "\n",
        "            # Don't compute if outside bandwidth\n",
        "            if not (abs(i - j) > bandwidth > 0):\n",
        "                a = math.exp((R[k, i + 1, j] - R[k, i, j] - D[k, i + 1, j]) * inv_gamma)\n",
        "                b = math.exp((R[k, i, j + 1] - R[k, i, j] - D[k, i, j + 1]) * inv_gamma)\n",
        "                c = math.exp((R[k, i + 1, j + 1] - R[k, i, j] - D[k, i + 1, j + 1]) * inv_gamma)\n",
        "                E[k, i, j] = E[k, i + 1, j] * a + E[k, i, j + 1] * b + E[k, i + 1, j + 1] * c\n",
        "\n",
        "        # Wait for other threads in this block\n",
        "        cuda.syncthreads()\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "class _SoftDTWCUDA(Function):\n",
        "    \"\"\"\n",
        "    CUDA implementation is inspired by the diagonal one proposed in https://ieeexplore.ieee.org/document/8400444:\n",
        "    \"Developing a pattern discovery method in time series data and its GPU acceleration\"\n",
        "    \"\"\"\n",
        "\n",
        "    @staticmethod\n",
        "    def forward(ctx, D, gamma, bandwidth):\n",
        "        dev = D.device\n",
        "        dtype = D.dtype\n",
        "        gamma = torch.cuda.FloatTensor([gamma])\n",
        "        bandwidth = torch.cuda.FloatTensor([bandwidth])\n",
        "\n",
        "        B = D.shape[0]\n",
        "        N = D.shape[1]\n",
        "        M = D.shape[2]\n",
        "        threads_per_block = max(N, M)\n",
        "        n_passes = 2 * threads_per_block - 1\n",
        "\n",
        "        # Prepare the output array\n",
        "        R = torch.ones((B, N + 2, M + 2), device=dev, dtype=dtype) * math.inf\n",
        "        R[:, 0, 0] = 0\n",
        "\n",
        "        # Run the CUDA kernel.\n",
        "        # Set CUDA's grid size to be equal to the batch size (every CUDA block processes one sample pair)\n",
        "        # Set the CUDA block size to be equal to the length of the longer sequence (equal to the size of the largest diagonal)\n",
        "        compute_softdtw_cuda[B, threads_per_block](cuda.as_cuda_array(D.detach()),\n",
        "                                                   gamma.item(), bandwidth.item(), N, M, n_passes,\n",
        "                                                   cuda.as_cuda_array(R))\n",
        "        ctx.save_for_backward(D, R.clone(), gamma, bandwidth)\n",
        "        return R[:, -2, -2]\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_output):\n",
        "        dev = grad_output.device\n",
        "        dtype = grad_output.dtype\n",
        "        D, R, gamma, bandwidth = ctx.saved_tensors\n",
        "\n",
        "        B = D.shape[0]\n",
        "        N = D.shape[1]\n",
        "        M = D.shape[2]\n",
        "        threads_per_block = max(N, M)\n",
        "        n_passes = 2 * threads_per_block - 1\n",
        "\n",
        "        D_ = torch.zeros((B, N + 2, M + 2), dtype=dtype, device=dev)\n",
        "        D_[:, 1:N + 1, 1:M + 1] = D\n",
        "\n",
        "        R[:, :, -1] = -math.inf\n",
        "        R[:, -1, :] = -math.inf\n",
        "        R[:, -1, -1] = R[:, -2, -2]\n",
        "\n",
        "        E = torch.zeros((B, N + 2, M + 2), dtype=dtype, device=dev)\n",
        "        E[:, -1, -1] = 1\n",
        "\n",
        "        # Grid and block sizes are set same as done above for the forward() call\n",
        "        compute_softdtw_backward_cuda[B, threads_per_block](cuda.as_cuda_array(D_),\n",
        "                                                            cuda.as_cuda_array(R),\n",
        "                                                            1.0 / gamma.item(), bandwidth.item(), N, M, n_passes,\n",
        "                                                            cuda.as_cuda_array(E))\n",
        "        E = E[:, 1:N + 1, 1:M + 1]\n",
        "        return grad_output.view(-1, 1, 1).expand_as(E) * E, None, None\n",
        "\n",
        "@jit(nopython=True)\n",
        "def compute_softdtw(D, gamma, bandwidth):\n",
        "    B = D.shape[0]\n",
        "    N = D.shape[1]\n",
        "    M = D.shape[2]\n",
        "    R = np.ones((B, N + 2, M + 2)) * np.inf\n",
        "    R[:, 0, 0] = 0\n",
        "    for b in range(B):\n",
        "        for j in range(1, M + 1):\n",
        "            for i in range(1, N + 1):\n",
        "\n",
        "                # Check the pruning condition\n",
        "                if 0 < bandwidth < np.abs(i - j):\n",
        "                    continue\n",
        "\n",
        "                r0 = -R[b, i - 1, j - 1] / gamma\n",
        "                r1 = -R[b, i - 1, j] / gamma\n",
        "                r2 = -R[b, i, j - 1] / gamma\n",
        "                rmax = max(max(r0, r1), r2)\n",
        "                rsum = np.exp(r0 - rmax) + np.exp(r1 - rmax) + np.exp(r2 - rmax)\n",
        "                softmin = - gamma * (np.log(rsum) + rmax)\n",
        "                R[b, i, j] = D[b, i - 1, j - 1] + softmin\n",
        "    return R\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "@jit(nopython=True)\n",
        "def compute_softdtw_backward(D_, R, gamma, bandwidth):\n",
        "    B = D_.shape[0]\n",
        "    N = D_.shape[1]\n",
        "    M = D_.shape[2]\n",
        "    D = np.zeros((B, N + 2, M + 2))\n",
        "    E = np.zeros((B, N + 2, M + 2))\n",
        "    D[:, 1:N + 1, 1:M + 1] = D_\n",
        "    E[:, -1, -1] = 1\n",
        "    R[:, :, -1] = -np.inf\n",
        "    R[:, -1, :] = -np.inf\n",
        "    R[:, -1, -1] = R[:, -2, -2]\n",
        "    for k in range(B):\n",
        "        for j in range(M, 0, -1):\n",
        "            for i in range(N, 0, -1):\n",
        "\n",
        "                if np.isinf(R[k, i, j]):\n",
        "                    R[k, i, j] = -np.inf\n",
        "\n",
        "                # Check the pruning condition\n",
        "                if 0 < bandwidth < np.abs(i - j):\n",
        "                    continue\n",
        "\n",
        "                a0 = (R[k, i + 1, j] - R[k, i, j] - D[k, i + 1, j]) / gamma\n",
        "                b0 = (R[k, i, j + 1] - R[k, i, j] - D[k, i, j + 1]) / gamma\n",
        "                c0 = (R[k, i + 1, j + 1] - R[k, i, j] - D[k, i + 1, j + 1]) / gamma\n",
        "                a = np.exp(a0)\n",
        "                b = np.exp(b0)\n",
        "                c = np.exp(c0)\n",
        "                E[k, i, j] = E[k, i + 1, j] * a + E[k, i, j + 1] * b + E[k, i + 1, j + 1] * c\n",
        "    return E[:, 1:N + 1, 1:M + 1]\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "class _SoftDTW(Function):\n",
        "    \"\"\"\n",
        "    CPU implementation based on https://github.com/Sleepwalking/pytorch-softdtw\n",
        "    \"\"\"\n",
        "\n",
        "    @staticmethod\n",
        "    def forward(ctx, D, gamma, bandwidth):\n",
        "        dev = D.device\n",
        "        dtype = D.dtype\n",
        "        gamma = torch.Tensor([gamma]).to(dev).type(dtype)  # dtype fixed\n",
        "        bandwidth = torch.Tensor([bandwidth]).to(dev).type(dtype)\n",
        "        D_ = D.detach().cpu().numpy()\n",
        "        g_ = gamma.item()\n",
        "        b_ = bandwidth.item()\n",
        "        R = torch.Tensor(compute_softdtw(D_, g_, b_)).to(dev).type(dtype)\n",
        "        ctx.save_for_backward(D, R, gamma, bandwidth)\n",
        "        return R[:, -2, -2]\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_output):\n",
        "        dev = grad_output.device\n",
        "        dtype = grad_output.dtype\n",
        "        D, R, gamma, bandwidth = ctx.saved_tensors\n",
        "        D_ = D.detach().cpu().numpy()\n",
        "        R_ = R.detach().cpu().numpy()\n",
        "        g_ = gamma.item()\n",
        "        b_ = bandwidth.item()\n",
        "        E = torch.Tensor(compute_softdtw_backward(D_, R_, g_, b_)).to(dev).type(dtype)\n",
        "        return grad_output.view(-1, 1, 1).expand_as(E) * E, None, None\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "class SoftDTW(torch.nn.Module):\n",
        "    \"\"\"\n",
        "    The soft DTW implementation that optionally supports CUDA\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, use_cuda, gamma=1.0, normalize=False, bandwidth=None, dist_func=None):\n",
        "        \"\"\"\n",
        "        Initializes a new instance using the supplied parameters\n",
        "        :param use_cuda: Flag indicating whether the CUDA implementation should be used\n",
        "        :param gamma: sDTW's gamma parameter\n",
        "        :param normalize: Flag indicating whether to perform normalization\n",
        "                          (as discussed in https://github.com/mblondel/soft-dtw/issues/10#issuecomment-383564790)\n",
        "        :param bandwidth: Sakoe-Chiba bandwidth for pruning. Passing 'None' will disable pruning.\n",
        "        :param dist_func: Optional point-wise distance function to use. If 'None', then a default Euclidean distance function will be used.\n",
        "        \"\"\"\n",
        "        super(SoftDTW, self).__init__()\n",
        "        self.normalize = normalize\n",
        "        self.gamma = gamma\n",
        "        self.bandwidth = 0 if bandwidth is None else float(bandwidth)\n",
        "        self.use_cuda = use_cuda\n",
        "\n",
        "        # Set the distance function\n",
        "        if dist_func is not None:\n",
        "            self.dist_func = dist_func\n",
        "        else:\n",
        "            self.dist_func = SoftDTW._euclidean_dist_func\n",
        "\n",
        "    def _get_func_dtw(self, x, y):\n",
        "        \"\"\"\n",
        "        Checks the inputs and selects the proper implementation to use.\n",
        "        \"\"\"\n",
        "        bx, lx, dx = x.shape\n",
        "        by, ly, dy = y.shape\n",
        "        # Make sure the dimensions match\n",
        "        assert bx == by  # Equal batch sizes\n",
        "        assert dx == dy  # Equal feature dimensions\n",
        "\n",
        "        use_cuda = self.use_cuda\n",
        "\n",
        "        if use_cuda and (lx > 1024 or ly > 1024):  # We should be able to spawn enough threads in CUDA\n",
        "                print(\"SoftDTW: Cannot use CUDA because the sequence length > 1024 (the maximum block size supported by CUDA)\")\n",
        "                use_cuda = False\n",
        "\n",
        "        # Finally, return the correct function\n",
        "        return _SoftDTWCUDA.apply if use_cuda else _SoftDTW.apply\n",
        "\n",
        "    @staticmethod\n",
        "    def _euclidean_dist_func(x, y):\n",
        "        \"\"\"\n",
        "        Calculates the Euclidean distance between each element in x and y per timestep\n",
        "        \"\"\"\n",
        "        n = x.size(1)\n",
        "        m = y.size(1)\n",
        "        d = x.size(2)\n",
        "        x = x.unsqueeze(2).expand(-1, n, m, d)\n",
        "        y = y.unsqueeze(1).expand(-1, n, m, d)\n",
        "        return torch.pow(x - y, 2).sum(3)\n",
        "\n",
        "    def forward(self, X, Y):\n",
        "        \"\"\"\n",
        "        Compute the soft-DTW value between X and Y\n",
        "        :param X: One batch of examples, batch_size x seq_len x dims\n",
        "        :param Y: The other batch of examples, batch_size x seq_len x dims\n",
        "        :return: The computed results\n",
        "        \"\"\"\n",
        "\n",
        "        # Check the inputs and get the correct implementation\n",
        "        func_dtw = self._get_func_dtw(X, Y)\n",
        "\n",
        "        if self.normalize:\n",
        "            # Stack everything up and run\n",
        "            x = torch.cat([X, X, Y])\n",
        "            y = torch.cat([Y, X, Y])\n",
        "            D = self.dist_func(x, y)\n",
        "            out = func_dtw(D, self.gamma, self.bandwidth)\n",
        "            out_xy, out_xx, out_yy = torch.split(out, X.shape[0])\n",
        "            return out_xy - 1 / 2 * (out_xx + out_yy)\n",
        "        else:\n",
        "            D_xy = self.dist_func(X, Y)\n",
        "            return func_dtw(D_xy, self.gamma, self.bandwidth)\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "def timed_run(a, b, sdtw):\n",
        "    \"\"\"\n",
        "    Runs a and b through sdtw, and times the forward and backward passes.\n",
        "    Assumes that a requires gradients.\n",
        "    :return: timing, forward result, backward result\n",
        "    \"\"\"\n",
        "    from timeit import default_timer as timer\n",
        "\n",
        "    # Forward pass\n",
        "    start = timer()\n",
        "    forward = sdtw(a, b)\n",
        "    end = timer()\n",
        "    t = end - start\n",
        "\n",
        "    grad_outputs = torch.ones_like(forward)\n",
        "\n",
        "    # Backward\n",
        "    start = timer()\n",
        "    grads = torch.autograd.grad(forward, a, grad_outputs=grad_outputs)[0]\n",
        "    end = timer()\n",
        "\n",
        "    # Total time\n",
        "    t += end - start\n",
        "\n",
        "    return t, forward, grads\n",
        "\n",
        "# ----------------------------------------------------------------------------------------------------------------------\n",
        "def profile(batch_size, seq_len_a, seq_len_b, dims, tol_backward):\n",
        "    sdtw = SoftDTW(False, gamma=1.0, normalize=False)\n",
        "    sdtw_cuda = SoftDTW(True, gamma=1.0, normalize=False)\n",
        "    n_iters = 6\n",
        "\n",
        "    print(\"Profiling forward() + backward() times for batch_size={}, seq_len_a={}, seq_len_b={}, dims={}...\".format(batch_size, seq_len_a, seq_len_b, dims))\n",
        "\n",
        "    times_cpu = []\n",
        "    times_gpu = []\n",
        "\n",
        "    for i in range(n_iters):\n",
        "        a_cpu = torch.rand((batch_size, seq_len_a, dims), requires_grad=True)\n",
        "        b_cpu = torch.rand((batch_size, seq_len_b, dims))\n",
        "        a_gpu = a_cpu.cuda()\n",
        "        b_gpu = b_cpu.cuda()\n",
        "\n",
        "        # GPU\n",
        "        t_gpu, forward_gpu, backward_gpu = timed_run(a_gpu, b_gpu, sdtw_cuda)\n",
        "\n",
        "        # CPU\n",
        "        t_cpu, forward_cpu, backward_cpu = timed_run(a_cpu, b_cpu, sdtw)\n",
        "\n",
        "        # Verify the results\n",
        "        assert torch.allclose(forward_cpu, forward_gpu.cpu())\n",
        "        assert torch.allclose(backward_cpu, backward_gpu.cpu(), atol=tol_backward)\n",
        "\n",
        "        if i > 0:  # Ignore the first time we run, in case this is a cold start (because timings are off at a cold start of the script)\n",
        "            times_cpu += [t_cpu]\n",
        "            times_gpu += [t_gpu]\n",
        "\n",
        "    # Average and log\n",
        "    avg_cpu = np.mean(times_cpu)\n",
        "    avg_gpu = np.mean(times_gpu)\n",
        "    print(\"\\tCPU:     \", avg_cpu)\n",
        "    print(\"\\tGPU:     \", avg_gpu)\n",
        "    print(\"\\tSpeedup: \", avg_cpu / avg_gpu)\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cnJpe8urvVYI"
      },
      "source": [
        "Models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "apYvRoDFoRwJ"
      },
      "outputs": [],
      "source": [
        "# all imports here\n",
        "import math\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import TensorDataset, DataLoader\n",
        "from torch.optim.lr_scheduler import _LRScheduler\n",
        "\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm\n",
        "import sklearn\n",
        "from copy import deepcopy\n",
        "\n",
        "class LSTM(nn.Module):\n",
        "\n",
        "    def __init__(self, num_classes, input_size, hidden_size, num_layers, bidirectional = False):\n",
        "        super(LSTM, self).__init__()\n",
        "        \n",
        "        self.num_classes = num_classes\n",
        "        self.num_layers = num_layers\n",
        "        self.input_size = input_size\n",
        "        self.hidden_size = hidden_size\n",
        "        self.seq_length = SEQ_LENGTH\n",
        "        self.bidrectional = bidirectional\n",
        "        \n",
        "        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,\n",
        "                            num_layers=num_layers, batch_first=True, bidirectional = bidirectional)\n",
        "        \n",
        "        self.fc = nn.Linear(hidden_size, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        h_0 = Variable(torch.zeros(\n",
        "            self.num_layers, x.size(0), self.hidden_size)).cuda()\n",
        "        \n",
        "        c_0 = Variable(torch.zeros(\n",
        "            self.num_layers, x.size(0), self.hidden_size)).cuda()\n",
        "        # Propagate input through LSTM\n",
        "        ula, (h_out, _) = self.lstm(x, (h_0, c_0))\n",
        "        \n",
        "        #h_out = h_out.view(-1, self.hidden_size)\n",
        "        out = self.fc(ula)\n",
        "        \n",
        "        return out\n",
        "\n",
        "import torch.nn as nn\n",
        "import math\n",
        "\n",
        "device = 'cuda'\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    '''Multi-head self-attention module'''\n",
        "    def __init__(self, D, H):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        self.H = H # number of heads\n",
        "        self.D = D # dimension\n",
        "        \n",
        "        self.wq = nn.Linear(D, D*H)\n",
        "        self.wk = nn.Linear(D, D*H)\n",
        "        self.wv = nn.Linear(D, D*H)\n",
        "\n",
        "        self.dense = nn.Linear(D*H, D)\n",
        "\n",
        "    def concat_heads(self, x):\n",
        "        '''(B, H, S, D) => (B, S, D*H)'''\n",
        "        B, H, S, D = x.shape\n",
        "        x = x.permute((0, 2, 1, 3)).contiguous()  # (B, S, H, D)\n",
        "        x = x.reshape((B, S, H*D))   # (B, S, D*H)\n",
        "        return x\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        '''(B, S, D*H) => (B, H, S, D)'''\n",
        "        B, S, D_H = x.shape\n",
        "        x = x.reshape(B, S, self.H, self.D)    # (B, S, H, D)\n",
        "        x = x.permute((0, 2, 1, 3))  # (B, H, S, D)\n",
        "        return x\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "\n",
        "        q = self.wq(x)  # (B, S, D*H)\n",
        "        k = self.wk(x)  # (B, S, D*H)\n",
        "        v = self.wv(x)  # (B, S, D*H)\n",
        "\n",
        "        q = self.split_heads(q)  # (B, H, S, D)\n",
        "        k = self.split_heads(k)  # (B, H, S, D)\n",
        "        v = self.split_heads(v)  # (B, H, S, D)\n",
        "\n",
        "        attention_scores = torch.matmul(q, k.transpose(-1, -2)) #(B,H,S,S)\n",
        "        attention_scores = attention_scores / math.sqrt(self.D)\n",
        "\n",
        "        # add the mask to the scaled tensor.\n",
        "        if mask is not None:\n",
        "            attention_scores += (mask * -1e9)\n",
        "        \n",
        "        attention_weights = nn.Softmax(dim=-1)(attention_scores)\n",
        "        scaled_attention = torch.matmul(attention_weights, v)  # (B, H, S, D)\n",
        "        concat_attention = self.concat_heads(scaled_attention) # (B, S, D*H)\n",
        "        output = self.dense(concat_attention)  # (B, S, D)\n",
        "\n",
        "        return output, attention_weights\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    '''Multi-head self-attention module'''\n",
        "    def __init__(self, D, H):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        self.H = H # number of heads\n",
        "        self.D = D # dimension\n",
        "        \n",
        "        self.wq = nn.Linear(D, D*H)\n",
        "        self.wk = nn.Linear(D, D*H)\n",
        "        self.wv = nn.Linear(D, D*H)\n",
        "\n",
        "        self.dense = nn.Linear(D*H, D)\n",
        "\n",
        "    def concat_heads(self, x):\n",
        "        '''(B, H, S, D) => (B, S, D*H)'''\n",
        "        B, H, S, D = x.shape\n",
        "        x = x.permute((0, 2, 1, 3)).contiguous()  # (B, S, H, D)\n",
        "        x = x.reshape((B, S, H*D))   # (B, S, D*H)\n",
        "        return x\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        '''(B, S, D*H) => (B, H, S, D)'''\n",
        "        B, S, D_H = x.shape\n",
        "        x = x.reshape(B, S, self.H, self.D)    # (B, S, H, D)\n",
        "        x = x.permute((0, 2, 1, 3))  # (B, H, S, D)\n",
        "        return x\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "\n",
        "        q = self.wq(x)  # (B, S, D*H)\n",
        "        k = self.wk(x)  # (B, S, D*H)\n",
        "        v = self.wv(x)  # (B, S, D*H)\n",
        "\n",
        "        q = self.split_heads(q)  # (B, H, S, D)\n",
        "        k = self.split_heads(k)  # (B, H, S, D)\n",
        "        v = self.split_heads(v)  # (B, H, S, D)\n",
        "\n",
        "        attention_scores = torch.matmul(q, k.transpose(-1, -2)) #(B,H,S,S)\n",
        "        attention_scores = attention_scores / math.sqrt(self.D)\n",
        "\n",
        "        # add the mask to the scaled tensor.\n",
        "        if mask is not None:\n",
        "            attention_scores += (mask * -1e9)\n",
        "        \n",
        "        attention_weights = nn.Softmax(dim=-1)(attention_scores)\n",
        "        scaled_attention = torch.matmul(attention_weights, v)  # (B, H, S, D)\n",
        "        concat_attention = self.concat_heads(scaled_attention) # (B, S, D*H)\n",
        "        output = self.dense(concat_attention)  # (B, S, D)\n",
        "\n",
        "        return output, attention_weights\n",
        "\n",
        "class MultiHeadAttentionCosformerNew(nn.Module):\n",
        "    '''Multi-head self-attention module'''\n",
        "    def __init__(self, D, H):\n",
        "        super(MultiHeadAttentionCosformerNew, self).__init__()\n",
        "        self.H = H # number of heads\n",
        "        self.D = D # dimension\n",
        "        \n",
        "        self.wq = nn.Linear(D, D*H)\n",
        "        self.wk = nn.Linear(D, D*H)\n",
        "        self.wv = nn.Linear(D, D*H)\n",
        "\n",
        "        self.dense = nn.Linear(D*H, D)\n",
        "\n",
        "    def concat_heads(self, x):\n",
        "        '''(B, H, S, D) => (B, S, D*H)'''\n",
        "        B, H, S, D = x.shape\n",
        "        x = x.permute((0, 2, 1, 3)).contiguous()  # (B, S, H, D)\n",
        "        x = x.reshape((B, S, H*D))   # (B, S, D*H)\n",
        "        return x\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        '''(B, S, D*H) => (B, H, S, D)'''\n",
        "        B, S, D_H = x.shape\n",
        "        x = x.reshape(B, S, self.H, self.D)    # (B, S, H, D)\n",
        "        x = x.permute((0, 2, 1, 3))  # (B, H, S, D)\n",
        "        return x\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "\n",
        "        q = self.wq(x)  # (B, S, D*H)\n",
        "        k = self.wk(x)  # (B, S, D*H)\n",
        "        v = self.wv(x)  # (B, S, D*H)\n",
        "\n",
        "        q = self.split_heads(q).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        k = self.split_heads(k).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        v = self.split_heads(v).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        B = q.shape[0]\n",
        "        S = q.shape[1]\n",
        "\n",
        "        q = torch.nn.functional.elu(q) + 1 # Sigmoid torch.nn.ReLU()\n",
        "        k = torch.nn.functional.elu(k) + 1 # Sigmoid torch.nn.ReLU()\n",
        "\n",
        "        # q, k, v -> [batch_size, seq_len, n_heads, d_head]\n",
        "        cos = (torch.cos(1.57*torch.arange(S)/S).unsqueeze(0)).repeat(B,1).cuda()\n",
        "        sin = (torch.sin(1.57*torch.arange(S)/S).unsqueeze(0)).repeat(B,1).cuda()\n",
        "        # cos, sin -> [batch_size, seq_len]\n",
        "        q_cos = torch.einsum('bsnd,bs->bsnd', q, cos)\n",
        "        q_sin = torch.einsum('bsnd,bs->bsnd', q, sin)\n",
        "        k_cos = torch.einsum('bsnd,bs->bsnd', k, cos)\n",
        "        k_sin = torch.einsum('bsnd,bs->bsnd', k, sin)\n",
        "        # q_cos, q_sin, k_cos, k_sin -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        kv_cos = torch.einsum('bsnx,bsnz->bnxz', k_cos, v)\n",
        "        # kv_cos -> [batch_size, n_heads, d_head, d_head]\n",
        "        qkv_cos = torch.einsum('bsnx,bnxz->bsnz', q_cos, kv_cos)\n",
        "        # qkv_cos -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        kv_sin = torch.einsum('bsnx,bsnz->bnxz', k_sin, v)\n",
        "        # kv_sin -> [batch_size, n_heads, d_head, d_head]\n",
        "        qkv_sin = torch.einsum('bsnx,bnxz->bsnz', q_sin, kv_sin)\n",
        "        # qkv_sin -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        # denominator\n",
        "        denominator = 1.0 / (torch.einsum('bsnd,bnd->bsn', q_cos, k_cos.sum(axis=1))\n",
        "                             + torch.einsum('bsnd,bnd->bsn',\n",
        "                                            q_sin, k_sin.sum(axis=1))\n",
        "                             + 1e-5)\n",
        "        # denominator -> [batch_size, seq_len, n_heads]\n",
        "\n",
        "        O = torch.einsum('bsnz,bsn->bsnz', qkv_cos +\n",
        "                              qkv_sin, denominator).contiguous()\n",
        "        # output -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        concat_attention = self.concat_heads(O.permute(0,2,1,3)) # (B, S, D*H)\n",
        "        output = self.dense(concat_attention)  # (B, S, D)\n",
        "\n",
        "        return output, None\n",
        "\n",
        "class MultiHeadAttentionCosSquareformerNew(nn.Module):\n",
        "    '''Multi-head self-attention module'''\n",
        "    def __init__(self, D, H):\n",
        "        super(MultiHeadAttentionCosSquareformerNew, self).__init__()\n",
        "        self.H = H # number of heads\n",
        "        self.D = D # dimension\n",
        "        \n",
        "        self.wq = nn.Linear(D, D*H)\n",
        "        self.wk = nn.Linear(D, D*H)\n",
        "        self.wv = nn.Linear(D, D*H)\n",
        "\n",
        "        self.dense = nn.Linear(D*H, D)\n",
        "\n",
        "    def concat_heads(self, x):\n",
        "        '''(B, H, S, D) => (B, S, D*H)'''\n",
        "        B, H, S, D = x.shape\n",
        "        x = x.permute((0, 2, 1, 3)).contiguous()  # (B, S, H, D)\n",
        "        x = x.reshape((B, S, H*D))   # (B, S, D*H)\n",
        "        return x\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        '''(B, S, D*H) => (B, H, S, D)'''\n",
        "        B, S, D_H = x.shape\n",
        "        x = x.reshape(B, S, self.H, self.D)    # (B, S, H, D)\n",
        "        x = x.permute((0, 2, 1, 3))  # (B, H, S, D)\n",
        "        return x\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "\n",
        "        q = self.wq(x)  # (B, S, D*H)\n",
        "        k = self.wk(x)  # (B, S, D*H)\n",
        "        v = self.wv(x)  # (B, S, D*H)\n",
        "\n",
        "        q = self.split_heads(q).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        k = self.split_heads(k).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        v = self.split_heads(v).permute(0,2,1,3)  # (B, S, H, D)\n",
        "        B = q.shape[0]\n",
        "        S = q.shape[1]\n",
        "\n",
        "        q = torch.nn.functional.elu(q) + 1 # Sigmoid torch.nn.ReLU()\n",
        "        k = torch.nn.functional.elu(k) + 1 # Sigmoid torch.nn.ReLU()\n",
        "\n",
        "        # q, k, v -> [batch_size, seq_len, n_heads, d_head]\n",
        "        cos = (torch.cos(3.1415*torch.arange(S)/S).unsqueeze(0)).repeat(B,1).cuda()\n",
        "        sin = (torch.sin(3.1415*torch.arange(S)/S).unsqueeze(0)).repeat(B,1).cuda()\n",
        "        # cos, sin -> [batch_size, seq_len]\n",
        "        q_cos = torch.einsum('bsnd,bs->bsnd', q, cos)\n",
        "        q_sin = torch.einsum('bsnd,bs->bsnd', q, sin)\n",
        "        k_cos = torch.einsum('bsnd,bs->bsnd', k, cos)\n",
        "        k_sin = torch.einsum('bsnd,bs->bsnd', k, sin)\n",
        "        # q_cos, q_sin, k_cos, k_sin -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        kv_cos = torch.einsum('bsnx,bsnz->bnxz', k_cos, v)\n",
        "        # kv_cos -> [batch_size, n_heads, d_head, d_head]\n",
        "        qkv_cos = torch.einsum('bsnx,bnxz->bsnz', q_cos, kv_cos)\n",
        "        # qkv_cos -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        kv_sin = torch.einsum('bsnx,bsnz->bnxz', k_sin, v)\n",
        "        # kv_sin -> [batch_size, n_heads, d_head, d_head]\n",
        "        qkv_sin = torch.einsum('bsnx,bnxz->bsnz', q_sin, kv_sin)\n",
        "        # qkv_sin -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        kv = torch.einsum('bsnx,bsnz->bnxz', k, v)\n",
        "        # kv -> [batch_size, n_heads, d_head, d_head]\n",
        "        qkv = torch.einsum('bsnx,bnxz->bsnz', q, kv)\n",
        "        # qkv_cos -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        # denominator\n",
        "        denominator = 1.0 / (torch.einsum('bsnd,bnd->bsn', q, k.sum(axis=1)) + torch.einsum('bsnd,bnd->bsn', q_cos, k_cos.sum(axis=1))\n",
        "                             + torch.einsum('bsnd,bnd->bsn',\n",
        "                                            q_sin, k_sin.sum(axis=1))\n",
        "                             + 1e-5)\n",
        "        # denominator -> [batch_size, seq_len, n_heads]\n",
        "\n",
        "        O = torch.einsum('bsnz,bsn->bsnz', qkv + qkv_cos +\n",
        "                              qkv_sin, denominator).contiguous()\n",
        "        # output -> [batch_size, seq_len, n_heads, d_head]\n",
        "\n",
        "        concat_attention = self.concat_heads(O.permute(0,2,1,3)) # (B, S, D*H)\n",
        "        output = self.dense(concat_attention)  # (B, S, D)\n",
        "\n",
        "        return output, None\n",
        "\n",
        "# Positional encodings\n",
        "def get_angles(pos, i, D):\n",
        "    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(D))\n",
        "    return pos * angle_rates\n",
        "\n",
        "\n",
        "def positional_encoding(D, position=20, dim=3, device=device):\n",
        "    angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                            np.arange(D)[np.newaxis, :],\n",
        "                            D)\n",
        "    # apply sin to even indices in the array; 2i\n",
        "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "    # apply cos to odd indices in the array; 2i+1\n",
        "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "    if dim == 3:\n",
        "        pos_encoding = angle_rads[np.newaxis, ...]\n",
        "    elif dim == 4:\n",
        "        pos_encoding = angle_rads[np.newaxis,np.newaxis,  ...]\n",
        "    return torch.tensor(pos_encoding, device=device)\n",
        "\n",
        "class TransformerLayer(nn.Module):\n",
        "    def __init__(self, D, H, hidden_mlp_dim, dropout_rate, attention_type='cosine_square'):\n",
        "        super(TransformerLayer, self).__init__()\n",
        "        self.dropout_rate = dropout_rate\n",
        "        self.mlp_hidden = nn.Linear(D, hidden_mlp_dim)\n",
        "        self.mlp_out = nn.Linear(hidden_mlp_dim, D)\n",
        "        self.layernorm1 = nn.LayerNorm(D, eps=1e-9)\n",
        "        self.layernorm2 = nn.LayerNorm(D, eps=1e-9)\n",
        "        self.dropout1 = nn.Dropout(dropout_rate)\n",
        "        self.dropout2 = nn.Dropout(dropout_rate)\n",
        "\n",
        "        if attention_type == 'cosine':\n",
        "          self.mha = MultiHeadAttentionCosformerNew(D, H)\n",
        "        elif attention_type == 'cosine_square':\n",
        "          self.mha = MultiHeadAttentionCosSquareformerNew(D, H)\n",
        "        else:\n",
        "          self.mha = MultiHeadAttention(D,H)\n",
        "\n",
        "    def forward(self, x, look_ahead_mask):\n",
        "        \n",
        "        attn, attn_weights = self.mha(x, look_ahead_mask)  # (B, S, D)\n",
        "        attn = self.dropout1(attn) # (B,S,D)\n",
        "        attn = self.layernorm1(attn + x) # (B,S,D)\n",
        "\n",
        "        mlp_act = torch.relu(self.mlp_hidden(attn))\n",
        "        mlp_act = self.mlp_out(mlp_act)\n",
        "        mlp_act = self.dropout2(mlp_act)\n",
        "        \n",
        "        output = self.layernorm2(mlp_act + attn)  # (B, S, D)\n",
        "\n",
        "        return output, attn_weights\n",
        "  \n",
        "class Transformer(nn.Module):\n",
        "    '''Transformer Decoder Implementating several Decoder Layers.\n",
        "    '''\n",
        "    def __init__(self, num_layers, D, H, hidden_mlp_dim, inp_features, out_features, dropout_rate, attention_type='cosine_square'):\n",
        "        super(Transformer, self).__init__()\n",
        "        self.attention_type = attention_type\n",
        "        self.sqrt_D = torch.tensor(math.sqrt(D))\n",
        "        self.num_layers = num_layers\n",
        "        self.input_projection = nn.Linear(inp_features, D) # multivariate input\n",
        "        self.output_projection = nn.Linear(D, out_features) # multivariate output\n",
        "        self.pos_encoding = positional_encoding(D)\n",
        "        self.dec_layers = nn.ModuleList([TransformerLayer(D, H, hidden_mlp_dim, \n",
        "                                        dropout_rate=dropout_rate, attention_type=self.attention_type\n",
        "                                       ) for _ in range(num_layers)])\n",
        "        self.dropout = nn.Dropout(dropout_rate)\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "        B, S, D = x.shape\n",
        "        # attention_weights = {}\n",
        "        x = self.input_projection(x)\n",
        "        x *= self.sqrt_D\n",
        "        \n",
        "        x += self.pos_encoding[:, :S, :]\n",
        "\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        for i in range(self.num_layers):\n",
        "            x, _ = self.dec_layers[i](x=x,\n",
        "                                          look_ahead_mask=mask)\n",
        "            # attention_weights['decoder_layer{}'.format(i + 1)] = block\n",
        "        \n",
        "        x = self.output_projection(x)\n",
        "        \n",
        "        return x, None # attention_weights # (B,S,S)\n",
        "\n",
        "class TransLSTM(nn.Module):\n",
        "    '''Transformer Decoder Implementating several Decoder Layers.\n",
        "    '''\n",
        "    def __init__(self, num_layers, D, H, hidden_mlp_dim, inp_features, out_features, dropout_rate, LSTM_module, attention_type='regular'):\n",
        "        super(TransLSTM, self).__init__()\n",
        "        self.attention_type = attention_type\n",
        "        self.sqrt_D = torch.tensor(math.sqrt(D))\n",
        "        self.num_layers = num_layers\n",
        "        self.input_projection = nn.Linear(inp_features, D) # multivariate input\n",
        "        self.output_projection = nn.Linear(D, 4) # multivariate output\n",
        "        self.fc = nn.Linear(4*2, out_features)\n",
        "        self.pos_encoding = positional_encoding(D)\n",
        "        self.dec_layers = nn.ModuleList([TransformerLayer(D, H, hidden_mlp_dim, \n",
        "                                        dropout_rate=dropout_rate, attention_type=self.attention_type\n",
        "                                       ) for _ in range(num_layers)])\n",
        "        self.dropout = nn.Dropout(dropout_rate)\n",
        "        self.LSTM = LSTM_module\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "        x_l = self.LSTM(x)\n",
        "        B, S, D = x.shape\n",
        "        attention_weights = {}\n",
        "        x = self.input_projection(x)\n",
        "        x *= self.sqrt_D\n",
        "        \n",
        "        x += self.pos_encoding[:, :S, :]\n",
        "\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        for i in range(self.num_layers):\n",
        "            x, block = self.dec_layers[i](x=x,\n",
        "                                          look_ahead_mask=mask)\n",
        "            attention_weights['decoder_layer{}'.format(i + 1)] = block\n",
        "        \n",
        "        x = self.output_projection(x)\n",
        "\n",
        "        x = torch.cat((x,x_l),axis=2)\n",
        "\n",
        "        x = self.fc(x)\n",
        "        \n",
        "        return x, attention_weights # (B,S,S)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jUVSeMq_vX94"
      },
      "source": [
        "Main"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VX6bRDGJ-wFU"
      },
      "outputs": [],
      "source": [
        "# all imports here\n",
        "import math\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import TensorDataset, DataLoader\n",
        "from torch.optim.lr_scheduler import _LRScheduler\n",
        "\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm\n",
        "import sklearn\n",
        "from copy import deepcopy\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "# function that implement the look_ahead mask for masking future time steps. \n",
        "def create_look_ahead_mask(size, device=device):\n",
        "    mask = torch.ones((size, size), device=device)\n",
        "    mask = torch.triu(mask, diagonal=1)\n",
        "    return mask  # (size, size)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "\n",
        "  dtw_loss = SoftDTW(use_cuda=True, gamma=0.1)\n",
        "  lmbda = 0.5\n",
        "\n",
        "  for SELECTED_COLUMN in [\"pm25_median\"]: # [\"pm25_median\", \"so2_median\", \"pm10_median\", \"no2_median\", \"o3_median\", \"co_median\", \"so2_median\"]:\n",
        "      \n",
        "      train_data = CityDataP(SELECTED_COLUMN, \"train\")\n",
        "      val_data = CityDataP(SELECTED_COLUMN, \"test\")\n",
        "\n",
        "      sampleLoader = DataLoader(train_data, 32, shuffle=True, num_workers=4)\n",
        "      val_loader = DataLoader(val_data, 4096, shuffle=False, num_workers=4)\n",
        "\n",
        "      lr = 0.001\n",
        "      n_epochs = 10\n",
        "\n",
        "      criterion = nn.MSELoss()\n",
        "      model = Transformer(num_layers=6, D=16, H=10, hidden_mlp_dim=32, inp_features=11, out_features=1, dropout_rate=0.1, attention_type='regular').to(device) # cosine_square, cosine, regular # 6L, 12H\n",
        "      # model = TransLSTM(num_layers=3, D=16, H=5, hidden_mlp_dim=32, inp_features=11, out_features=1, dropout_rate=0.2, LSTM_module = LSTM(4, INPUT_DIM+1, HIDDEN_DIM, LAYER_DIM, bidirectional = False).to(device), attention_type='regular').to(device) # cosine_square, cosine, regular # 6L, 12H\n",
        "      # model = LSTM(1, INPUT_DIM+1, HIDDEN_DIM, LAYER_DIM).cuda()\n",
        "\n",
        "      opt = torch.optim.Adam(model.parameters(), lr=lr)\n",
        "      \n",
        "      print('Start model training')\n",
        "      best_mse = 2000.0\n",
        "      best_model = None\n",
        "\n",
        "      for epoch in range(1, n_epochs + 1):\n",
        "          epoch_loss = 0\n",
        "          batch_idx = 0\n",
        "          bar = tqdm(sampleLoader)\n",
        "\n",
        "          model.train()\n",
        "          for x_batch, y_batch, _ in bar:\n",
        "              model.train()\n",
        "              x_batch = x_batch.cuda().float()\n",
        "              y_batch = y_batch.cuda().float()\n",
        "\n",
        "              mask = create_look_ahead_mask(x_batch.shape[1])\n",
        "              out, _ = model(x_batch, mask)\n",
        "              opt.zero_grad()\n",
        "\n",
        "              loss = criterion(out[:,-1,:], y_batch[:,-1,:]) + lmbda * dtw_loss(out.cuda(),y_batch.cuda()).mean()\n",
        "\n",
        "              epoch_loss = (epoch_loss*batch_idx + loss.item())/(batch_idx+1)\n",
        "              loss.backward()\n",
        "              opt.step()\n",
        "\n",
        "              bar.set_description(str(epoch_loss))\n",
        "              batch_idx += 1\n",
        "\n",
        "          # Evaluation\n",
        "          model.eval()\n",
        "          mse_list = []\n",
        "          total_se = 0.0\n",
        "          total_pe = 0.0\n",
        "          total_valid = 0.0\n",
        "\n",
        "          for x_val, _, y_val in val_loader:\n",
        "              x_val, y_val = [t.cuda().float() for t in (x_val, y_val)]\n",
        "              mask = create_look_ahead_mask(x_val.shape[1])\n",
        "\n",
        "              out, _ = model(x_val, mask)\n",
        "\n",
        "              ytrue = y_val[:,-1,:].squeeze().cpu().numpy()\n",
        "              ypred = out[:,-1,:].squeeze().cpu().detach().numpy()\n",
        "              \n",
        "              true_valid = np.isnan(ytrue) != 1\n",
        "              ytrue = ytrue[true_valid] #np.nan_to_num(ytrue, 0)\n",
        "              ypred = ypred[true_valid]\n",
        "\n",
        "              if normalization_type == 'mean_std':\n",
        "                ytrue = (ytrue * col_std[SELECTED_COLUMN]) + col_mean2[SELECTED_COLUMN]\n",
        "                ypred = (ypred * col_std[SELECTED_COLUMN]) + col_mean2[SELECTED_COLUMN]\n",
        "              \n",
        "              else:\n",
        "                ytrue = (ytrue * col_max[SELECTED_COLUMN])\n",
        "                ypred = (ypred * col_max[SELECTED_COLUMN])\n",
        "\n",
        "              se = (ytrue - ypred)**2 # np.square(ytrue - ypred)\n",
        "              pe = np.abs((ytrue - ypred) / (ytrue + 0.0001))\n",
        "              total_se += np.sum(se)\n",
        "              total_pe += np.sum(pe)\n",
        "              total_valid += np.sum(true_valid)\n",
        "\n",
        "          eval_mse = total_se / total_valid # np.mean(se) # \n",
        "          eval_mape = total_pe / total_valid # np.mean(pe) # \n",
        "          print('valid samples:', total_valid)\n",
        "          print('Eval MSE: ', eval_mse)\n",
        "          print('Eval RMSE: {}: '.format(SELECTED_COLUMN), np.sqrt(eval_mse))\n",
        "          print('Eval MAPE: {}: '.format(SELECTED_COLUMN), eval_mape*100)\n",
        "\n",
        "          if eval_mse < best_mse:\n",
        "            best_model = deepcopy(model)\n",
        "            best_mse = eval_mse\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "train.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}